package matchbox;

import java.util.ArrayList;
import java.util.Random;

import inference.GraphNode;
import inference.BeliefMessage;
import probability.GaussianDistribution;
import probability.Indicator;
import probability.ProbabilityDistribution;
import utils.Printer;
import basics.DenseMatrix;
import basics.DenseVector;
import basics.FeaturesSet;
import basics.Vector;
import basics.VectorMatrix;

public class STFactor extends MatchboxFactor {

	private double[] _means;
	private double[] _vars;

	private Vector _x;

	private GaussianDistribution _g = new GaussianDistribution();

	public STFactor(String name) {
		super(name, Type.Operation);
	}

	public void Preprocess(VectorMatrix x) {
		_x = (Vector) x;

		double kmean = 0;
		double kvar = 1;

		if (_means == null) {
			init(x.length());
		}

		for (int i = 0; i < x.length(); i++) {
			kmean += (x.get(i) * _means[i]); // x.get(i).timesCopy(_means[i]).sumAllElements();
			kvar += (x.get(i) * _vars[i]);
		}

		_g.setMu(kmean);
		_g.setSigma2(kvar);

		_count++;
		if (_xMean == null) {
			_xMean = (DenseVector) _x.copy();
		} else {
			_xMean.addEqual(_x.minus(_xMean).times(1. / (double) (_count)));
		}
	}

	private void init(int length) {
		_means = new double[length];
		_vars = new double[length];
		Random r = new Random();

		for (int i = 0; i < length; i++) {
			_means[i] = r.nextDouble() / 10.; // 0.; // 1. / (double) (length * lDimension);
			_vars[i] = 1.;
		}
	}

	@Override
	public BeliefMessage toCached(GraphNode next, String opt) {
		BeliefMessage m = null;
		if (next.is(NodeNames.S) || next.is(NodeNames.T)) {
			m = new BeliefMessage(_g); // new DenseVector(new double[] {
		} else {
			ProbabilityDistribution pd = edges(NodeNames.T).nodes().get(0).toCached(this).getDistribution();
			m = new BeliefMessage(_x);
		}

		return m;
	}

	@Override
	public BeliefMessage to(GraphNode next, String opt) {
		BeliefMessage m = null;
		if (next.is(NodeNames.S) || next.is(NodeNames.T) || next.is(NodeNames.B)) {
			m = new BeliefMessage(_g); // new DenseVector(new double[] {
		} else {
			ProbabilityDistribution pd = edges(NodeNames.T).nodes().get(0).to(this).getDistribution();
			// _PlatentVar.setParam(0, _g.getParam(0).addCopy(pd.getParam(0)));
			// _PlatentVar.setParam(0, _g.getParam(1).addCopy(pd.getParam(1)));
			// m = new BeliefMessage(_PlatentVar);
			m = new BeliefMessage(_x);
			// DenseVector x = new DenseVector(new double[] { _PlatentVar });
		}
		// _g.measure(next.to(this)) });
		Printer.print("F_ST>" + next + ":" + m.getDistribution().toString());

		return m;
	}

	private ArrayList<GaussianDistribution> _messagesIn = new ArrayList<GaussianDistribution>();

	public void updateLatent(GaussianDistribution g) {
		// _messagesIn.add(g);
		// if (_messagesIn.size() >= 1) {
		// g = GaussianDistribution.product(_messagesIn);
		// trueskillUpdate(g);
		// _messagesIn.clear();
		// }

		for (int i = 0; i < _means.length; i++) {
			GaussianDistribution t = new GaussianDistribution(g.getMu() - (_g.getMu() - _means[i]), g.getSigma2()
					- (_g.getSigma2() - _vars[i]));
			t = (GaussianDistribution) t.times(new GaussianDistribution(_means[i], _vars[i]));
			_means[i] = t.getMu();
			_vars[i] = t.getSigma2();
		}
	}

	private void trueskillUpdate(GaussianDistribution g) {
		// ArrayList<GaussianDistribution> fTov = new ArrayList<GaussianDistribution>();
		ArrayList<GaussianDistribution> vTof = new ArrayList<GaussianDistribution>();
		double[] coeff = new double[_means.length];
		double[] m = new double[_means.length];
		double[] v = new double[_means.length];

		for (int i = 0; i < _x.length(); i++) {
			if (_x.get(i) == 0)
				continue;
			for (int k = 0; k < _x.length(); k++) {
				if (k == i) {
					vTof.add(g);
					coeff[k] = 1.;
					// fTov.add(_g);
				} else {
					vTof.add(new GaussianDistribution(_means[k], _vars[k]));
					// fTov.add(new GaussianDistribution(1. / _means.length, 1));
					coeff[k] = (-1.) * _x.get(k);
				}
				// double denom = == 0 ? 0.001 : _x.get(i);
				coeff[k] = _x.get(i) == 0 ? 0 : coeff[k] * (1. / _x.get(i));
			}

			GaussianDistribution res = Indicator.resolveIndicator(vTof, coeff); // fTov,
			m[i] = res.getMu();
			v[i] = res.getSigma2();
			vTof.clear();
			// fTov.clear();
		}

		for (int i = 0; i < _means.length; i++) {
			if (_x.get(i) == 0)
				continue;
			GaussianDistribution prod = GaussianDistribution.product(new GaussianDistribution(_means[i], _vars[i]),
					new GaussianDistribution(m[i], v[i]));
			_means[i] = prod.getMu();
			_vars[i] = prod.getSigma2();
		}
	}

	private static void updateMeanVars(GaussianDistribution gd, Vector x, double[] mean, double[] var, int length) {
		for (int i = 0; i < length; i++) {
			// GaussianDistribution temp = GaussianDistribution.product(g, new GaussianDistribution(_means[i],
			// _vars[i]));
			// mean[i] = temp.getMu();
			// var[i] = temp.getSigma2();

			// _means[i] = m - (_g.getMu() - _means[i] * _x.get(i)) / _x.get(i);
			// _vars[i] = v - (_g.getMu() - _vars[i] * _x.get(i)) / _x.get(i);
		}
		Printer.print("mean: ", mean);
		Printer.print("Var: ", var);
	}

	public GaussianDistribution getPuv() {
		GaussianDistribution g = null;
		for (int i = 0; i < _means.length; i++) {
			GaussianDistribution t = new GaussianDistribution(_means[i], _vars[i]);
			if (g == null) {
				g = t;
			} else {
				GaussianDistribution.product(g, t);
			}
		}
		return g;
	}

	public Vector getX() {
		return _x;
	}

	public double[] getMeans() {
		return _means;
	}

	private int _count = 0;
	private DenseVector _xMean;

	public DenseVector getXMean() {
		return _xMean;
	}

	@Override
	public GaussianDistribution ep(GraphNode next) {
		GaussianDistribution pd = null;
		// if (next == null) {
		// pd = (GaussianDistribution) toCached(null).getDistribution();
		// pd = pd.divide(_g);
		// updateLatent(pd);
		// _g =
		// }

		return pd;
	}

}
